{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FhGuhbZ6M5tl"
      },
      "source": [
        "##### Copyright 2022 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "AwOEIRJC6Une"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EIdT9iu_Z4Rb"
      },
      "source": [
        "# Matrix approximation with Core APIs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bBIlTPscrIT9"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/core/matrix_core\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/core/matrix_core.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/core/matrix_core.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/core/matrix_core.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGw8TF2vtzru"
      },
      "source": [
        "## Introduction \n",
        "\n",
        "This notebook uses the [TensorFlow Core low-level APIs](https://www.tensorflow.org/guide/core) to showcase TensorFlow's capabilities as a high-performance scientific computing platform. Visit the [Core APIs overview](https://www.tensorflow.org/guide/core) to learn more about TensorFlow Core and its intended use cases.\n",
        "\n",
        "This tutorial explores the technique of [singular value decomposition](https://developers.google.com/machine-learning/recommendation/collaborative/matrix) (SVD) and its applications for low-rank approximation problems. The SVD is used to factorize real or complex matrices and has a variety of use cases in data science such as image compression. The images for this tutorial come from Google Brain's [Imagen](https://imagen.research.google/) project. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5_FdwaovEkCC"
      },
      "source": [
        ">![svd_intro](http://tensorflow.org/images/core/svd_intro.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nchsZfwEVtVs"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1rRo8oNqZ-Rj"
      },
      "outputs": [],
      "source": [
        "import matplotlib\n",
        "from matplotlib.image import imread\n",
        "from matplotlib import pyplot as plt\n",
        "import requests\n",
        "# Preset Matplotlib figure sizes.\n",
        "matplotlib.rcParams['figure.figsize'] = [16, 9]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9xQKvCJ85kCQ"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "print(tf.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "so_ewq3gAoEI"
      },
      "source": [
        "## SVD fundamentals\n",
        "\n",
        "The singular value decomposition of a matrix, ${\\mathrm{A}}$, is determined by the following factorization:\n",
        "\n",
        "$${\\mathrm{A}} = {\\mathrm{U}} \\Sigma {\\mathrm{V}}^T$$\n",
        "\n",
        "where\n",
        "\n",
        "* $\\underset{m \\times n}{\\mathrm{A}}$: input matrix where $m \\geq n$\n",
        "* $\\underset{m \\times n}{\\mathrm{U}}$: orthogonal matrix, ${\\mathrm{U}}^T{\\mathrm{U}} = {\\mathrm{I}}$, with each column, $u_i$, denoting a left singular vector of ${\\mathrm{A}}$\n",
        "* $\\underset{n \\times n}{\\Sigma}$: diagonal matrix with each diagonal entry, $\\sigma_i$, denoting a singular value of ${\\mathrm{A}}$\n",
        "* $\\underset{n \\times n}{{\\mathrm{V}}^T}$: orthogonal matrix, ${\\mathrm{V}}^T{\\mathrm{V}} = {\\mathrm{I}}$, with each row, $v_i$, denoting a right singular vector of ${\\mathrm{A}}$\n",
        "\n",
        "When $m < n$, ${\\mathrm{U}}$ and $\\Sigma$ both have dimension $(m \\times m)$, and ${\\mathrm{V}}^T$ has dimension $(m \\times n)$."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "enGGGXCQKNv8"
      },
      "source": [
        ">![svd_full](http://tensorflow.org/images/core/svd_full.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NlP-cBdSKLtc"
      },
      "source": [
        "TensorFlow's linear algebra package has a function, `tf.linalg.svd`, which can be used to compute the singular value decomposition of one or more matrices. Start by defining a simple matrix and computing its SVD factorization.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C3QAcgyoeIpv"
      },
      "outputs": [],
      "source": [
        "A = tf.random.uniform(shape=[40,30])\n",
        "# Compute the SVD factorization\n",
        "s, U, V = tf.linalg.svd(A)\n",
        "# Define Sigma and V Transpose\n",
        "S = tf.linalg.diag(s)\n",
        "V_T = tf.transpose(V)\n",
        "# Reconstruct the original matrix\n",
        "A_svd = U@S@V_T\n",
        "# Visualize \n",
        "plt.bar(range(len(s)), s);\n",
        "plt.xlabel(\"Singular value rank\")\n",
        "plt.ylabel(\"Singular value\")\n",
        "plt.title(\"Bar graph of singular values\");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6H_C9WhFACm4"
      },
      "source": [
        "The `tf.einsum` function can be used to directly compute the matrix reconstruction from the outputs of `tf.linalg.svd`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TPE6QeMtADUn"
      },
      "outputs": [],
      "source": [
        "A_svd = tf.einsum('s,us,vs -> uv',s,U,V)\n",
        "print('\\nReconstructed Matrix, A_svd', A_svd)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x1m6JIsM9DLP"
      },
      "source": [
        "## Low rank approximation with the SVD\n",
        "\n",
        "The rank of a  matrix, ${\\mathrm{A}}$, is determined by the dimension of the vector space spanned by its columns. \n",
        "The SVD can be used to approximate a matrix with a lower rank, which ultimately decreases the dimensionality of data required to store the information represented by the matrix.\n",
        "\n",
        "The rank-r approximation of ${\\mathrm{A}}$ in terms of the SVD is defined by the formula:\n",
        "\n",
        "$${\\mathrm{A_r}} = {\\mathrm{U_r}} \\Sigma_r {\\mathrm{V_r}}^T$$\n",
        "\n",
        "where\n",
        "\n",
        "* $\\underset{m \\times r}{\\mathrm{U_r}}$: matrix consisting of the first $r$ columns of ${\\mathrm{U}}$\n",
        "* $\\underset{r \\times r}{\\Sigma_r}$:  diagonal matrix consisting of the first  $r$ singular values in $\\Sigma$\n",
        "* $\\underset{r \\times n}{\\mathrm{V_r}}^T$: matrix consisting of the first $r$ rows of ${\\mathrm{V}}^T$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nJWMJu36QyUV"
      },
      "source": [
        ">![svd_approx](http://tensorflow.org/images/core/svd_approx.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TkiVUxeaQybq"
      },
      "source": [
        "Start by writing a function to compute the rank-r approximation of a given matrix. This low-rank approximation procedure is used for image compression; therefore, it is also helpful to compute the physical data sizes for each approximation. For simplicity, assume that data size for an rank-r approximated matrix is equal to the total number of elements required to compute the approximation. Next, write a function to visualize the original matrix, $\\mathrm{A}$ its rank-r approximation, $\\mathrm{A}_r$ and the error matrix, $|\\mathrm{A} - \\mathrm{A}_r|$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2oY3pMPagJrO"
      },
      "outputs": [],
      "source": [
        "def rank_r_approx(s, U, V, r, verbose=False):\n",
        "  # Compute the matrices necessary for a rank-r approximation\n",
        "  s_r, U_r, V_r = s[..., :r], U[..., :, :r], V[..., :, :r] # ... implies any number of extra batch axes\n",
        "  # Compute the low-rank approximation and its size\n",
        "  A_r = tf.einsum('...s,...us,...vs->...uv',s_r,U_r,V_r)\n",
        "  A_r_size = tf.size(U_r) + tf.size(s_r) + tf.size(V_r)\n",
        "  if verbose:\n",
        "    print(f\"Approximation Size: {A_r_size}\")\n",
        "  return A_r, A_r_size\n",
        "\n",
        "def viz_approx(A, A_r):\n",
        "  # Plot A, A_r, and A - A_r\n",
        "  vmin, vmax = 0, tf.reduce_max(A)\n",
        "  fig, ax = plt.subplots(1,3)\n",
        "  mats = [A, A_r, abs(A - A_r)]\n",
        "  titles = ['Original A', 'Approximated A_r', 'Error |A - A_r|']\n",
        "  for i, (mat, title) in enumerate(zip(mats, titles)):\n",
        "    ax[i].pcolormesh(mat, vmin=vmin, vmax=vmax)\n",
        "    ax[i].set_title(title)\n",
        "    ax[i].axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O3ZRkYCkX2FQ"
      },
      "outputs": [],
      "source": [
        "print(f\"Original Size of A: {tf.size(A)}\")\n",
        "s, U, V = tf.linalg.svd(A)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S1DR83VMX4cM"
      },
      "outputs": [],
      "source": [
        "# Rank-15 approximation\n",
        "A_15, A_15_size = rank_r_approx(s, U, V, 15, verbose = True)\n",
        "viz_approx(A, A_15)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KgFT70XFX57E"
      },
      "outputs": [],
      "source": [
        "# Rank-3 approximation\n",
        "A_3, A_3_size = rank_r_approx(s, U, V, 3, verbose = True)\n",
        "viz_approx(A, A_3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DS4XoSlTJgX0"
      },
      "source": [
        "As expected, using lower ranks results in less-accurate approximations. However, the quality of these low-rank approximations are often good enough in real world scenarios. Also note that the main goal of low-rank approximation with SVD \n",
        "is to reduce the dimensionality of the data but not to reduce the disk space of the data itself. However, as the input matrices become higher-dimensional, many low-rank approximations also end up benefiting from reduced data size. This reduction benefit is why the process is applicable for image compression problems."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IhsaiOnnZs6M"
      },
      "source": [
        "## Image loading\n",
        "\n",
        "The following image is available on the [Imagen](https://imagen.research.google/) home page. Imagen is a text-to-image diffusion model developed by Google Research's Brain team. An AI created this image based on the prompt: \"A photo of a Corgi dog riding a bike in Times Square. It is wearing sunglasses and a beach hat.\" How cool is that! You can also change the url below to any .jpg link to load in a custom image of choice. \n",
        "\n",
        "Start by reading in and visualizing the image. After reading a JPEG file, Matplotlib outputs a matrix, ${\\mathrm{I}}$, of shape $(m \\times n \\times 3)$ which represents a 2-dimensional image with 3 color channels for red, green and blue respectively."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OVsZOQUAZ2C7"
      },
      "outputs": [],
      "source": [
        "img_link = \"https://imagen.research.google/main_gallery_images/a-photo-of-a-corgi-dog-riding-a-bike-in-times-square.jpg\"\n",
        "img_path = requests.get(img_link, stream=True).raw\n",
        "I = imread(img_path, 0)\n",
        "print(\"Input Image Shape:\", I.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qvs7uftcZ54x"
      },
      "outputs": [],
      "source": [
        "def show_img(I):\n",
        "  # Display the image in matplotlib\n",
        "  img = plt.imshow(I)\n",
        "  plt.axis('off')\n",
        "  return"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZbesXO3HZ6Qs"
      },
      "outputs": [],
      "source": [
        "show_img(I)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tdnUBVg_JoOa"
      },
      "source": [
        "## The image compression algorithm\n",
        "\n",
        "Now, use the SVD to compute low-rank approximations of the sample image. Recall that the image is of shape $(1024 \\times 1024 \\times 3)$ and that the theory SVD only applies for 2-dimensional matrices. This means that the sample image has to be batched into 3 equal-size matrices that correspond to each of the 3 color channels. This can be done so by transposing the matrix to be of shape $(3 \\times 1024 \\times 1024)$. In order to clearly visualize the approximation error, rescale the RGB values of the image from $[0,255]$ to $[0,1]$. Remember to clip the approximated values to fall within this interval before visualizing them. The `tf.clip_by_value` function is useful for this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i7DDp0h7oSIk"
      },
      "outputs": [],
      "source": [
        "def compress_image(I, r, verbose=False):\n",
        "  # Compress an image with the SVD given a rank \n",
        "  I_size = tf.size(I)\n",
        "  print(f\"Original size of image: {I_size}\")\n",
        "  # Compute SVD of image\n",
        "  I = tf.convert_to_tensor(I)/255\n",
        "  I_batched = tf.transpose(I, [2, 0, 1]) # einops.rearrange(I, 'h w c -> c h w')\n",
        "  s, U, V = tf.linalg.svd(I_batched)\n",
        "  # Compute low-rank approximation of image across each RGB channel\n",
        "  I_r, I_r_size = rank_r_approx(s, U, V, r)\n",
        "  I_r = tf.transpose(I_r, [1, 2, 0]) # einops.rearrange(I_r, 'c h w -> h w c')\n",
        "  I_r_prop = (I_r_size / I_size)\n",
        "  if verbose:\n",
        "    # Display compressed image and attributes\n",
        "    print(f\"Number of singular values used in compression: {r}\")\n",
        "    print(f\"Compressed image size: {I_r_size}\")\n",
        "    print(f\"Proportion of original size: {I_r_prop:.3f}\")\n",
        "    ax_1 = plt.subplot(1,2,1)\n",
        "    show_img(tf.clip_by_value(I_r,0.,1.))\n",
        "    ax_1.set_title(\"Approximated image\")\n",
        "    ax_2 = plt.subplot(1,2,2)\n",
        "    show_img(tf.clip_by_value(0.5+abs(I-I_r),0.,1.))\n",
        "    ax_2.set_title(\"Error\")\n",
        "  return I_r, I_r_prop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RGQ_rTyKDX9F"
      },
      "source": [
        "Now, compute rank-r approximations for the following ranks : 100, 50, 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7GlKkVLGDjre"
      },
      "outputs": [],
      "source": [
        "I_100, I_100_prop = compress_image(I, 100, verbose=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XdvUkF5_E75D"
      },
      "outputs": [],
      "source": [
        "I_50, I_50_prop = compress_image(I, 50, verbose=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MsCNZ8416Sbk"
      },
      "outputs": [],
      "source": [
        "I_10, I_10_prop = compress_image(I, 10, verbose=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RfYYBhcuNkvH"
      },
      "source": [
        "## Evaluating approximations\n",
        "\n",
        "There are a variety of interesting methods to measure the effectiveness and have more control over matrix approximations."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D2Lotde9Zg7v"
      },
      "source": [
        "### Compression factor vs rank\n",
        "\n",
        "For each of the above approximations, observe how the data sizes change with the rank."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O1ariNQe6Wbl"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(11,6))\n",
        "plt.plot([100, 50, 10], [I_100_prop, I_50_prop, I_10_prop])\n",
        "plt.xlabel(\"Rank\")\n",
        "plt.ylabel(\"Proportion of original image size\")\n",
        "plt.title(\"Compression factor vs rank\");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dvHcLRj2QoDg"
      },
      "source": [
        "Based on this plot, there is a linear relationship between an approximated image's compression factor and its rank. To explore this further, recall that the data size of an approximated matrix, ${\\mathrm{A}}_r$, is defined as the total number of elements required for its computation. The following equations can be used to find the relationship between compression factor and rank:\n",
        "\n",
        "$$x = (m \\times r) + r + (r \\times n) = r \\times (m + n + 1)$$\n",
        "\n",
        "$$c = \\large \\frac{x}{y} = \\frac{r \\times (m + n + 1)}{m \\times n}$$\n",
        "\n",
        "where\n",
        "\n",
        "* $x$: size of ${\\mathrm{A_r}}$\n",
        "* $y$: size of ${\\mathrm{A}}$\n",
        "* $c = \\frac{x}{y}$: compression factor\n",
        "* $r$: rank of the approximation\n",
        "* $m$ and $n$: row and column dimensions of ${\\mathrm{A}}$\n",
        "\n",
        "In order to find the rank, $r$, that is necessary to compress an image to a desired factor, $c$, the above equation can be rearranged to solve for $r$:\n",
        "\n",
        "$$r = ⌊{\\large\\frac{c \\times m \\times n}{m + n + 1}}⌋$$\n",
        "\n",
        "Note that this formula is independent of the color channel dimension since each of the RGB approximations do not affect each other. Now, write a function to compress an input image given a desired compression factor."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "viVO-I60QynI"
      },
      "outputs": [],
      "source": [
        "def compress_image_with_factor(I, compression_factor, verbose=False):\n",
        "  # Returns a compressed image based on a desired compression factor\n",
        "  m,n,o = I.shape\n",
        "  r = int((compression_factor * m * n)/(m + n + 1))\n",
        "  I_r, I_r_prop = compress_image(I, r, verbose=verbose)\n",
        "  return I_r"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gWSv58J6LSRQ"
      },
      "source": [
        "Compress an image to 15% of its original size."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HVeeloIwQ1b6"
      },
      "outputs": [],
      "source": [
        "compression_factor = 0.15\n",
        "I_r_img = compress_image_with_factor(I, compression_factor, verbose=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LkeRyms7jZMd"
      },
      "source": [
        "### Cumulative sum of singular values\n",
        "\n",
        "The cumulative sum of singular values can be a useful indicator for the amount of energy captured by a rank-r approximation. Visualize the RGB-averaged cumulative proportion of singular values in the sample image. The `tf.cumsum` function can be useful for this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CteJ6VbKlndu"
      },
      "outputs": [],
      "source": [
        "def viz_energy(I):\n",
        "  # Visualize the energy captured based on rank\n",
        "  # Computing SVD\n",
        "  I = tf.convert_to_tensor(I)/255\n",
        "  I_batched = tf.transpose(I, [2, 0, 1]) \n",
        "  s, U, V = tf.linalg.svd(I_batched)\n",
        "  # Plotting average proportion across RGB channels \n",
        "  props_rgb = tf.map_fn(lambda x: tf.cumsum(x)/tf.reduce_sum(x), s)\n",
        "  props_rgb_mean = tf.reduce_mean(props_rgb, axis=0)\n",
        "  plt.figure(figsize=(11,6))\n",
        "  plt.plot(range(len(I)), props_rgb_mean, color='k')\n",
        "  plt.xlabel(\"Rank / singular value number\")\n",
        "  plt.ylabel(\"Cumulative proportion of singular values\")\n",
        "  plt.title(\"RGB-averaged proportion of energy captured by the first 'r' singular values\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vl9PKow-GgCp"
      },
      "outputs": [],
      "source": [
        "viz_energy(I)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vQtwimKuQP19"
      },
      "source": [
        "It looks like over 90% of the energy in this image is captured within the first 100 singular values. Now, write a function to compress an input image given a desired energy retention factor."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fum5Cvm7R5vH"
      },
      "outputs": [],
      "source": [
        "def compress_image_with_energy(I, energy_factor, verbose=False):\n",
        "  # Returns a compressed image based on a desired energy factor\n",
        "  # Computing SVD\n",
        "  I_rescaled = tf.convert_to_tensor(I)/255\n",
        "  I_batched = tf.transpose(I_rescaled, [2, 0, 1]) \n",
        "  s, U, V = tf.linalg.svd(I_batched)\n",
        "  # Extracting singular values\n",
        "  props_rgb = tf.map_fn(lambda x: tf.cumsum(x)/tf.reduce_sum(x), s)\n",
        "  props_rgb_mean = tf.reduce_mean(props_rgb, axis=0)\n",
        "  # Find closest r that corresponds to the energy factor\n",
        "  r = tf.argmin(tf.abs(props_rgb_mean - energy_factor)) + 1\n",
        "  actual_ef = props_rgb_mean[r]\n",
        "  I_r, I_r_prop = compress_image(I, r, verbose=verbose)\n",
        "  print(f\"Proportion of energy captured by the first {r} singular values: {actual_ef:.3f}\")\n",
        "  return I_r"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y_rChG0OLby1"
      },
      "source": [
        "Compress an image to retain 75% of its energy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xDXBaZQ4c5jF"
      },
      "outputs": [],
      "source": [
        "energy_factor = 0.75\n",
        "I_r_img = compress_image_with_energy(I, energy_factor, verbose=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2tmqTW0CYX-v"
      },
      "source": [
        "### Error and singular values\n",
        "\n",
        "There is also an interesting relationship between the approximation error and the singular values. It turns out that the squared Frobenius norm of the approximation is equal to the sum of the squares of its singular values that were left out:\n",
        "\n",
        "$${||A - A_r||}^2 = \\sum_{i=r+1}^{R}σ_i^2$$\n",
        "\n",
        "Test out this relationship with a rank-10 approximation of the example matrix in the beginning of this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hctOvN8BckiS"
      },
      "outputs": [],
      "source": [
        "s, U, V = tf.linalg.svd(A)\n",
        "A_10, A_10_size = rank_r_approx(s, U, V, 10)\n",
        "squared_norm = tf.norm(A - A_10)**2\n",
        "s_squared_sum = tf.reduce_sum(s[10:]**2)\n",
        "print(f\"Squared Frobenius norm: {squared_norm:.3f}\")\n",
        "print(f\"Sum of squared singular values left out: {s_squared_sum:.3f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vgGQuV-yqYZH"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "This notebook introduced the process of implementing the singular value decomposition with TensorFlow and applying it to write an image compression algorithm. Here are a few more tips that may help:\n",
        "\n",
        "*   The [TensorFlow Core APIs](https://www.tensorflow.org/guide/core) can be utilized for a variety of high-performance scientific computing use cases.\n",
        "*   To learn more about TensorFlow's linear algebra functionalities, visit the docs for the [linalg module](https://www.tensorflow.org/api_docs/python/tf/linalg).\n",
        "*   The SVD can also be applied to build [recommendation systems](https://developers.google.com/machine-learning/recommendation/labs/movie-rec-programming-exercise).\n",
        "\n",
        "\n",
        "For more examples of using the TensorFlow Core APIs, check out the [guide](https://www.tensorflow.org/guide/core). If you want learn more about loading and preparing data, see the tutorials on [image data loading](https://www.tensorflow.org/tutorials/load_data/images) or [CSV data loading](https://www.tensorflow.org/tutorials/load_data/csv)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "matrix_core.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
